package com.jie.aspect;

import com.jie.annotation.Limit;
import com.jie.constants.RedisLimitType;
import com.jie.dto.UserDetailDTO;
import com.jie.exception.BizException;
import com.jie.utils.IpUtils;
import com.jie.utils.UserUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.util.Collections;
import java.util.Objects;


@Slf4j
@Aspect
@Component
public class LimitAspect {

    @Autowired
    private RedisTemplate redisTemplate;

    @Pointcut("@annotation(com.jie.annotation.Limit)")
    public void pointcut() {

    }

    @Around("pointcut()")
    public Object around(ProceedingJoinPoint point) throws Throwable {
        HttpServletRequest request = ((ServletRequestAttributes) Objects.requireNonNull(RequestContextHolder.getRequestAttributes())).getRequest();

        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        Limit limit = method.getAnnotation(Limit.class);
        RedisLimitType limitType = limit.limitType();
         //key前缀
        String prefix= limit.prefix();
        //key为IP和用户id
        String key="";
        //用户ip地址
        String ip = IpUtils.getIpAddress(request);
        // 令牌桶参数
        int capacity = limit.capacity(); // 桶的大小
        double rate = limit.rate(); // 每秒填充的令牌数
        int period = limit.period(); // 令牌桶的 refill 周期，单位为秒
        String name= limit.name();
        //判断根据什么类型进行限流
        switch (limitType) {
            case IP:
                key = prefix + ":" + ip;
                break;
            case PICTURE:
                //做的全局的限流
                key = prefix + ":" + name;
                break;
        }
        //执行令牌桶算法
        String luaScript = buildTokenBucketScript(capacity, rate, period);
        RedisScript<Long> redisScript = new DefaultRedisScript<>(luaScript, Long.class);
        Long result = (Long) redisTemplate.execute(redisScript,
                Collections.singletonList(key),
                rate, capacity, period, System.currentTimeMillis());
        if (result != null && result > 0) {
            log.info("当前访问的IP: {}", ip);
            return point.proceed();
        } else {
            log.warn("访问人数过多限制的IP: {}",ip);
            throw new BizException("请求过于频繁，请稍后再试");
        }
    }

    /**
     * 令牌桶限流lua
     * @param capacity 桶的大小
     * @param rate 每秒填充的令牌数
     * @param period 周期时间
     * @return
     */
    private String buildTokenBucketScript(int capacity, double rate, int period) {
        return "local key = KEYS[1] " + // 桶的key
                "local rate = tonumber(ARGV[1]) " + // 每秒填充的令牌数
                "local capacity = tonumber(ARGV[2]) " + // 桶的大小
                "local period = tonumber(ARGV[3]) " + // 周期时间（秒）
                "local now = tonumber(ARGV[4]) " + // 当前时间（秒）
                "local capacityName=key .. '' "+//桶的名称
                "local keyLastRefillTime = key .. ':lastRefillTime' " + // 构建lastRefillTime的key
                "local lastRefillTime = tonumber(redis.call('get', keyLastRefillTime)) " +
                "local currentTokens = tonumber(redis.call('get', key)) " +
                "if not lastRefillTime or not currentTokens then " + // 如果是首次访问，初始化lastRefillTime和桶
                "  lastRefillTime = now " +
                "  currentTokens = capacity " + // 初始化桶的令牌数为最大容量
                "  redis.call('set', key, currentTokens) " + // 设置桶的初始令牌数
                "  redis.call('set', keyLastRefillTime, lastRefillTime) " + // 设置lastRefillTime
                "end " +
                "local elapsed = now - lastRefillTime " + // 计算自上次填充以来经过的时间 这个是毫秒 需要换为秒
                "local tokensToAdd = rate * elapsed/1000 " + // 计算应添加的令牌数
                "tokensToAdd = math.floor(tokensToAdd) " + // 令牌数必须是整数
                "local newTokens = currentTokens + tokensToAdd " + // 计算新的令牌数
                "newTokens = math.min(capacity, newTokens) " + // 确保不超过桶的容量
                "redis.call('set', key, newTokens) " + // 更新桶中的令牌数
                "redis.call('set', keyLastRefillTime, now) " + // 更新lastRefillTime
                "if newTokens > 0 then " +
                "  redis.call('decrby', key, 1) " + // 从桶中减去一个令牌
                "  return 1 " + // 请求未被限流
                "else " +
                "  return 0 " + // 请求被限流
                "end";
    }
}

